import torch
from torch import nn
# --- 构建数据集的库 ---
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# 自定义的模型
from CustomNet import CustomNet

# --- 对于我们的数据而言，首先是有
# 一个共同的文件夹，其次是要有分类的文件夹 ---

# -----------------------------构建数据集-----------------------------------
# 对图像数据预处理
transform = transforms.Compose([
    transforms.ToTensor()  # 为了将数据转化为Pytorch对象（必要的内容）
])
# 加载图像，并且进行分类
train_datasets = datasets.ImageFolder("./data/numbers", transform=transform)
# 针对数据集进行分批次加工
train_loader = DataLoader(train_datasets, batch_size=1000, shuffle=True)
# -----------------------------构建模型-----------------------------------
model = CustomNet()
# -----------------------------构建训练参数-----------------------------------
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # 配置运行设备
# 损失函数
criterion = nn.CrossEntropyLoss()
# 优化器
sgd = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
# 训练次数
epochs = 10
# -----------------------------开始训练-----------------------------------
model = model.to(device)  # 设置模型设备
for epoch in range(epochs):
    losses = []
    for images, labels in iter(train_loader):
        sgd.zero_grad()
        predict_labels = model(images.to(device))
        loss = criterion(predict_labels, labels.to(device))
        loss.backward()
        sgd.step()

        losses.append(loss.item())  # 将每一个batch_size 统计起来

    total_loss = sum(losses) / len(losses)
    print(f"epoch: {epoch + 1} / {epochs} ---- loss: {total_loss:.4f}")
# -----------------------------保存训练好的模型-----------------------------------
torch.save(model, "numbers.pth")
